import os
import mysql.connector
import re
import json
from flask import Flask, request, jsonify, render_template
from cozepy import Coze, TokenAuth, Message, ChatStatus, COZE_CN_BASE_URL
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

app = Flask(__name__)

# --- Database Connection ---
def get_db_connection():
    """Establishes a connection to the MySQL database."""
    try:
        conn = mysql.connector.connect(
            host=os.getenv("DB_HOST"),
            user=os.getenv("DB_USER"),
            password=os.getenv("DB_PASSWORD"),
            database=os.getenv("DB_NAME")
        )
        return conn
    except mysql.connector.Error as err:
        print(f"Error connecting to database: {err}")
        return None

# --- Coze Client Setup ---
def create_coze_client():
    """Initializes the Coze client with detailed error handling."""
    try:
        coze_api_token = os.getenv("COZE_API_TOKEN")
        if not coze_api_token:
            raise ValueError("COZE_API_TOKEN environment variable is not set")
        
        bot_id = os.getenv("COZE_BOT_ID")
        if not bot_id:
            raise ValueError("COZE_BOT_ID environment variable is not set")
            
        user_id = os.getenv("COZE_USER_ID", "medical_app_user_001")
        
        print(f"Initializing Coze client with Bot ID: {bot_id}")
        
        # Initialize Coze client
        coze = Coze(
            auth=TokenAuth(token=coze_api_token), 
            base_url=COZE_CN_BASE_URL
        )
        
        # 测试连接
        print("Coze client initialized successfully")
        return coze, bot_id, user_id
        
    except Exception as e:
        print(f"Error initializing Coze client: {e}")
        return None, None, None

# Initialize Coze client
coze_client, coze_bot_id, coze_user_id = create_coze_client()

def clean_coze_response(text):
    """清理Coze响应中的JSON数据和无关内容，但保留完整的医疗建议"""
    if not text:
        return ""
    
    # 移除所有JSON对象和插件数据
    text = re.sub(r'\{[^{}]*"finish_reason"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"api_key"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"plugin_id"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"plugin_name"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"api_id"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"api_name"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"plugin_type"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"document"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"link"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"card"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"from_module"[^{}]*\}', '', text)
    text = re.sub(r'\{[^{}]*"from_unit"[^{}]*\}', '', text)
    text = re.sub(r'\{"data":\{[^{}]*\}\}', '', text)
    text = re.sub(r'\{"msg_type":.*?\}', '', text)
    
    # 移除常见的无关问题（保留医疗相关内容）
    unwanted_phrases = [
        '提供一些乳腺癌的治疗成功案例',
        '乳腺癌的发病率和死亡率是多少？',
        '上述肿瘤数据中的.*?是多少？',
        '提供一些预防乳腺癌的方法',
        '乳腺癌常见的治疗方法有哪些？',
        '乳腺癌的早期症状有哪些？',
        '这些数据具体反映了肿瘤的哪些特征？',
        '乳腺超声和钼靶检查的区别是什么？'
    ]
    
    for phrase in unwanted_phrases:
        text = re.sub(phrase, '', text)
    
    # 移除多余的空格和换行，但保留合理的段落结构
    text = re.sub(r'\n+', '\n', text)
    text = re.sub(r' +', ' ', text)
    text = text.strip()
    
    return text

def extract_result_from_response(response_text):
    """从Coze响应中提取判断结果和治疗建议 - 改进版"""
    cleaned_text = clean_coze_response(response_text)
    
    print(f"Cleaned response: {cleaned_text}")
    
    # 初始化默认值
    judgment_result = "无法确定"
    treatment_advice = cleaned_text
    
    # 更精确的模式匹配
    judgment_patterns = [
        r'判断结果[：:]\s*([^\n]+?(?=治疗建议|$))',
        r'判断结果[：:]\s*([^。]+?(?=治疗建议|$))',
        r'判断结果[：:]\s*(.+?(?=治疗建议|$))'
    ]
    
    treatment_patterns = [
        r'治疗建议[：:]\s*([^\n]+?(?=判断结果|给出治疗建议的依据|肿瘤的良恶性|提供一些预防|$))',
        r'治疗建议[：:]\s*([^。]+?(?=判断结果|给出治疗建议的依据|肿瘤的良恶性|提供一些预防|$))',
        r'治疗建议[：:]\s*(.+?(?=判断结果|给出治疗建议的依据|肿瘤的良恶性|提供一些预防|$))'
    ]
    
    # 查找判断结果
    for pattern in judgment_patterns:
        judgment_match = re.search(pattern, cleaned_text, re.DOTALL)
        if judgment_match:
            judgment_result = judgment_match.group(1).strip()
            # 清理判断结果中的额外内容
            judgment_result = re.split(r'[。\n]', judgment_result)[0]
            break
    
    # 查找治疗建议 - 使用更精确的匹配
    treatment_found = False
    for pattern in treatment_patterns:
        treatment_match = re.search(pattern, cleaned_text, re.DOTALL)
        if treatment_match:
            treatment_advice = treatment_match.group(1).strip()
            treatment_found = True
            break
    
    # 如果没找到格式化的治疗建议，尝试其他方法
    if not treatment_found:
        # 方法1: 找到判断结果后的第一个段落作为治疗建议
        judgment_index = cleaned_text.find('判断结果')
        if judgment_index != -1:
            # 找到判断结果行的结束位置
            judgment_line_end = cleaned_text.find('\n', judgment_index)
            if judgment_line_end != -1:
                remaining_text = cleaned_text[judgment_line_end:].strip()
                # 找到治疗建议的结束位置（在额外问题之前）
                end_markers = [
                    '给出治疗建议的依据',
                    '肿瘤的良恶性会随着时间变化吗？',
                    '提供一些预防乳腺癌的建议',
                    '这些数据具体反映了',
                    '乳腺癌有哪些常见的症状',
                    '乳腺超声和钼靶检查的区别'
                ]
                
                min_end_index = len(remaining_text)
                for marker in end_markers:
                    marker_index = remaining_text.find(marker)
                    if marker_index != -1 and marker_index < min_end_index:
                        min_end_index = marker_index
                
                if min_end_index < len(remaining_text):
                    treatment_advice = remaining_text[:min_end_index].strip()
                else:
                    treatment_advice = remaining_text
    
    # 最终清理治疗建议
    treatment_advice = treatment_advice.strip()
    
    # 如果治疗建议为空或与判断结果相同，使用默认值
    if not treatment_advice or treatment_advice == judgment_result:
        treatment_advice = "基于专业分析提供的医疗建议"
    
    # 确保治疗建议不以判断结果开头
    if treatment_advice.startswith(judgment_result):
        treatment_advice = treatment_advice[len(judgment_result):].strip()
        if treatment_advice.startswith('治疗建议：'):
            treatment_advice = treatment_advice[5:].strip()
    
    # 移除治疗建议中可能残留的额外问题
    extra_questions = [
        '给出治疗建议的依据',
        '肿瘤的良恶性会随着时间变化吗？',
        '提供一些预防乳腺癌的建议',
        '这些数据具体反映了肿瘤的哪些特征？',
        '乳腺癌有哪些常见的症状？',
        '乳腺超声和钼靶检查的区别是什么？'
    ]
    
    for question in extra_questions:
        if question in treatment_advice:
            treatment_advice = treatment_advice.split(question)[0].strip()
    
    return judgment_result, treatment_advice

# --- Flask Routes ---
@app.route('/')
def index():
    """Renders the main page."""
    return render_template('index.html')

@app.route('/health')
def health_check():
    """Health check endpoint to verify Coze connection."""
    if coze_client and coze_bot_id:
        return jsonify({
            "status": "healthy",
            "coze_client_initialized": coze_client is not None,
            "bot_id_set": coze_bot_id is not None
        })
    else:
        return jsonify({
            "status": "unhealthy",
            "error": "Coze client not properly initialized"
        }), 500

@app.route('/predict', methods=['POST'])
def predict():
    """Handles the prediction request."""
    try:
        # Check if Coze client is available
        if coze_client is None or coze_bot_id is None:
            return jsonify({"error": "AI服务暂时不可用，请检查系统配置"}), 500
        
        data = request.get_json()
        if not data:
            return jsonify({"error": "请求数据为空"}), 400
        
        # The list of expected features
        features = [
            "radius_mean", "texture_mean", "perimeter_mean", "area_mean",
            "smoothness_mean", "compactness_mean", "concavity_mean",
            "concave_points_mean", "symmetry_mean", "fractal_dimension_mean",
            "radius_se", "texture_se"
        ]

        # Validate that all features are present
        missing_features = [feature for feature in features if feature not in data]
        if missing_features:
            return jsonify({
                "error": "缺少必要的输入数据",
                "missing_features": missing_features
            }), 400

        # Format data for the LLM prompt
        data_input_str = "\n".join([f"- {feature}: {data[feature]}" for feature in features])
        
        # Create the prompt for Coze - 强调格式要求
        # 在predict函数中改进提示词
        prompt_text = f"""
你是一位顶尖且专业的乳腺癌医生，拥有深厚的专业知识与极为丰富的临床经验。凭借专业素养，你能精准剖析输入的数据，判断患者是否患有乳腺癌。若数据不符合医学现实情况，你要清晰指出不符合现实之处，并简要给出科学合理的治疗建议。

请分析以下肿瘤测量数据：
{data_input_str}

请严格按照以下格式回复，不要添加任何额外的问题、讨论或其他内容：

判断结果：<明确写出是否患有乳腺癌，例如"良性肿瘤"或"恶性肿瘤"或"数据不符合现实">
治疗建议：<简洁清晰地给出具有针对性的治疗建议或者详细说明数据不符合医学现实的具体情况。>

重要要求：
1. 只输出上述两行内容，不要添加任何其他文字
2. 治疗建议部分控制在200字以内
3. 不要提出任何问题或请求更多信息
4. 不要包含JSON、文档引用或其他技术数据
"""
        
        print(f"Sending request to Coze bot {coze_bot_id}")
        
        # Run the Coze chat
        try:
            chat_poll = coze_client.chat.create_and_poll(
                bot_id=coze_bot_id,
                user_id=coze_user_id,
                additional_messages=[
                    Message.build_user_question_text(prompt_text),
                ],
            )
            
            print(f"Coze response status: {chat_poll.chat.status}")
            print(f"Token usage: {chat_poll.chat.usage.token_count if chat_poll.chat.usage else 'N/A'}")
            
        except Exception as coze_error:
            print(f"Coze API error: {coze_error}")
            return jsonify({"error": "AI服务调用失败"}), 500
        
        # Check if chat completed successfully
        if chat_poll.chat.status != ChatStatus.COMPLETED:
            return jsonify({"error": "AI服务响应未完成"}), 500
        
        # Extract the response text from all messages
        response_text = ""
        for message in chat_poll.messages:
            if hasattr(message, 'content') and message.content:
                content_str = str(message.content)
                # 只添加非JSON内容
                if content_str and not content_str.strip().startswith('{'):
                    response_text += content_str + "\n"
        
        print(f"Raw Coze response: {response_text}")
        
        # 解析响应文本
        judgment_result, treatment_advice = extract_result_from_response(response_text)
        
        print(f"Parsed judgment: {judgment_result}")
        print(f"Parsed treatment: {treatment_advice}")
        
        # 确定数据库存储的预测结果
        if "恶性" in judgment_result:
            db_prediction = "Malignant Tumors"
        elif "良性" in judgment_result:
            db_prediction = "Benign Tumors"
        elif "不符合现实" in judgment_result:
            db_prediction = "Data Unrealistic"
        else:
            db_prediction = "Uncertain"
        
        # --- Store data in MySQL ---
        conn = get_db_connection()
        if conn:
            cursor = conn.cursor()
            
            # Prepare data for insertion
            insert_data = {feature: data.get(feature) for feature in features}
            insert_data['prediction_result'] = db_prediction

            # Use backticks for column names to avoid issues with reserved words or special chars
            columns = '`, `'.join(insert_data.keys())
            placeholders = ', '.join(['%s'] * len(insert_data))
            
            sql = f"INSERT INTO tumor_predictions (`{columns}`) VALUES ({placeholders})"
            
            try:
                cursor.execute(sql, list(insert_data.values()))
                conn.commit()
                print("Data successfully stored in database")
            except mysql.connector.Error as err:
                print(f"Database insertion error: {err}")
                conn.rollback()
            finally:
                cursor.close()
                conn.close()
        else:
            print("Warning: Database connection failed. Result not saved.")
        
        # 计算置信度
        confidence = 0
        if "恶性" in judgment_result or "良性" in judgment_result:
            confidence = 95
        elif "不符合现实" in judgment_result:
            confidence = 90
        else:
            confidence = 50
        
        return jsonify({
            "judgment_result": judgment_result,
            "treatment_advice": treatment_advice,
            "db_prediction": db_prediction,
            "confidence": confidence
        })

    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        return jsonify({"error": "服务器内部错误"}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)